import json
import os
import time
import glob
from datetime import datetime
from openai import OpenAI
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 常量定义
HISTORY_DIR = "chat_histories"
KNOWLEDGE_DIR = "knowledge_base"

class HistoryManager:
    """历史记录管理类，支持多文件加载"""
    
    def __init__(self, history_dir=HISTORY_DIR):
        self.history_dir = history_dir
        os.makedirs(history_dir, exist_ok=True)
    
    def list_histories(self, limit=10):
        """列出所有历史记录文件（按时间倒序）"""
        files = []
        for fname in os.listdir(self.history_dir):
            if fname.endswith(".json") and fname != "latest.json":
                file_path = os.path.join(self.history_dir, fname)
                files.append({
                    "id": len(files) + 1,
                    "filename": fname,
                    "path": file_path,
                    "mtime": os.path.getmtime(file_path)
                })
        
        # 按修改时间倒序排序
        files.sort(key=lambda x: x["mtime"], reverse=True)
        return files[:limit]
    
    def load_history(self, filename):
        """加载单个历史记录文件"""
        file_path = os.path.join(self.history_dir, filename)
        if os.path.exists(file_path):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                print(f"⚠️ 加载历史记录失败: {filename} - {e}")
        return None
    
    def load_multiple_histories(self, filenames):
        """加载多个历史记录文件并合并"""
        merged_history = []
        for filename in filenames:
            history = self.load_history(filename)
            if history:
                # 跳过重复的系统消息
                if merged_history and history[0].get("role") == "system":
                    merged_history.extend(history[1:])
                else:
                    merged_history.extend(history)
        return merged_history
    
    def save_history(self, messages, custom_name=None):
        """保存当前对话历史"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = custom_name or f"history_{timestamp}.json"
        filepath = os.path.join(self.history_dir, filename)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(messages, f, ensure_ascii=False, indent=4)
        
        # 更新最新记录
        latest_path = os.path.join(self.history_dir, "latest.json")
        with open(latest_path, 'w', encoding='utf-8') as f:
            json.dump(messages, f, ensure_ascii=False, indent=4)
        
        return filepath

class KnowledgeLoader:
    """知识库加载类"""
    
    def __init__(self, knowledge_dir=KNOWLEDGE_DIR):
        self.knowledge_dir = knowledge_dir
        os.makedirs(knowledge_dir, exist_ok=True)
    
    def load(self):
        """加载知识库内容"""
        content = []
        extensions = ('.txt', '.md', '.csv', '.json')
        
        for ext in extensions:
            pattern = os.path.join(self.knowledge_dir, '**', f'*{ext}')
            for file_path in glob.glob(pattern, recursive=True):
                if os.path.isfile(file_path):
                    try:
                        # 大文件只读取前10K字符
                        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                            text = f.read(10240) if os.path.getsize(file_path) > 102400 else f.read()
                            content.append({
                                "path": os.path.basename(file_path),
                                "content": text
                            })
                    except Exception as e:
                        print(f"⚠️ 文件读取失败: {file_path} - {e}")
        return content

def create_client():
    """创建OpenAI客户端"""
    try:
        return OpenAI(
            api_key=os.getenv("API_KEY"),
            base_url=os.getenv("API_URL")
        )
    except Exception as e:
        print(f"创建客户端时出错: {e}")
        return None

def stream_chat(client, messages):
    """流式聊天函数"""
    try:
        print("\nAI助手正在思考...\n" + "-" * 40)
        
        response = client.chat.completions.create(
            model="Qwen2.5-72B-Instruct",
            messages=messages,
            stream=True,
            temperature=0.7,
            max_tokens=2000
        )
        
        full_response = ""
        print("AI助手: ", end='', flush=True)
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                content = chunk.choices[0].delta.content
                print(content, end='', flush=True)
                full_response += content
                time.sleep(0.01)
        
        print("\n" + "-" * 40)
        return full_response
        
    except Exception as e:
        print(f"\nAPI调用错误: {e}")
        return None

def build_rag_prompt(user_query, knowledge, history):
    """构建RAG提示词"""
    # 系统指令
    system_msg = {
        "role": "system",
        "content": "你是一个专业的问答助手，请严格根据提供的知识内容回答问题。"
    }
    
    # 知识库上下文
    knowledge_context = "\n".join([f"【{k['path']}】\n{k['content'][:500]}" for k in knowledge])
    knowledge_msg = {
        "role": "user",
        "content": f"{knowledge_context}\n\n【当前问题】{user_query}"
    }
    
    # 组合历史记录
    return [system_msg] + history + [knowledge_msg]

def main():
    # 定义ANSI颜色代码
    ORANGE = '\033[38;5;214m'
    RESET = '\033[0m'
    
    print("=" * 60)
    print(ORANGE + "🤖 多历史文件加载聊天系统" + RESET)
    print("=" * 60)
    print("命令指南:")
    print(f"  {ORANGE}list{RESET}    - 列出历史对话")
    print(f"  {ORANGE}load N{RESET}  - 加载多个历史文件（如: load 1,3,5）")
    print(f"  {ORANGE}save{RESET}    - 保存当前对话")
    print(f"  {ORANGE}clear{RESET}   - 清空当前对话")
    print(f"  {ORANGE}reload{RESET}  - 重载知识库")
    print(f"  {ORANGE}exit{RESET}    - 退出程序")
    print(f"  {ORANGE}ctrl+C{RESET} - 保存对话并退出")
    print("=" * 60)
    
    client = create_client()
    if not client:
        print("无法创建客户端，程序退出")
        return
    
    # 初始化管理器
    history_manager = HistoryManager()
    knowledge_loader = KnowledgeLoader()
    
    # 加载知识库
    knowledge = knowledge_loader.load()
    print(f"\n📚 已加载 {len(knowledge)} 个知识文件")
    
    # 初始化对话
    messages = [
        {"role": "system", "content": "你是一个友好且专业的AI助手，请用中文回答用户问题"}
    ]
    
    # 尝试加载最新记录
    latest_history = history_manager.load_history("latest.json")
    if latest_history:
        messages = latest_history
        print("\n🔄 已恢复上次对话")
    
    while True:
        try:
            user_input = input("\n👤 您: ").strip()
            if not user_input:
                continue
                
            # 命令处理
            if user_input.lower() == 'exit':
                history_manager.save_history(messages)
                print("\n👋 对话已保存，再见！")
                break
                
            if user_input.lower() == 'list':
                files = history_manager.list_histories()
                if files:
                    print("\n📜 历史对话记录:")
                    for file in files:
                        print(f"  [{file['id']}] {file['filename']}")
                else:
                    print("\nℹ️ 无历史记录")
                continue
                
            if user_input.lower().startswith('load '):
                try:
                    # 解析多个ID (如: "1,3,5")
                    ids = [int(id.strip()) for id in user_input[5:].split(',')]
                    files = history_manager.list_histories()
                    
                    # 获取文件名
                    filenames = []
                    for id in ids:
                        if 1 <= id <= len(files):
                            filenames.append(files[id-1]["filename"])
                    
                    if filenames:
                        # 加载并合并多个历史文件
                        merged_history = history_manager.load_multiple_histories(filenames)
                        messages = merged_history
                        print(f"\n✅ 已加载 {len(filenames)} 个历史文件")
                    else:
                        print("\n❌ 未找到匹配的历史文件")
                except:
                    print("\n⚠️ 命令格式: load <编号1>,<编号2>,...")
                continue
                
            if user_input.lower() == 'save':
                saved_path = history_manager.save_history(messages)
                print(f"\n💾 对话已保存至: {saved_path}")
                continue
                
            if user_input.lower() == 'clear':
                messages = [messages[0]]
                print("\n🗑️ 对话已清空")
                continue
                
            if user_input.lower() == 'reload':
                knowledge = knowledge_loader.load()
                print(f"\n🔄 已重载 {len(knowledge)} 个知识文件")
                continue
            
            # 添加用户消息
            messages.append({"role": "user", "content": user_input})
            
            # 构建RAG提示并获取回复
            rag_prompt = build_rag_prompt(user_input, knowledge, messages)
            response = stream_chat(client, rag_prompt)
            
            if response:
                # 添加AI回复
                messages.append({"role": "assistant", "content": response})
                
                # 保持上下文长度
                if len(messages) > 50:
                    messages = [messages[0]] + messages[-49:]
            else:
                print("回复获取失败，请重试")
                messages.pop()  # 移除失败的用户消息
        
        except KeyboardInterrupt:
            history_manager.save_history(messages)
            print("\n\n⏹️ 对话已保存，程序中断")
            break
        except Exception as e:
            print(f"\n⚠️ 系统错误: {e}")

if __name__ == "__main__":
    main()